#!/usr/bin/env python3
#encoding=UTF8
import json
import sys
import time
from pymongo import MongoClient, errors

def usage():
    print(sys.argv[0]+" (discover|rs)")
    sys.exit(1)

if len(sys.argv) < 1:
    usage()

TIMEOUT_REQUEST = 2

ACTION = sys.argv[1]


class MongoDB(object):
    """main script class"""
    # pylint: disable=too-many-instance-attributes
    def __init__(self):
        self.mongo_host = "127.0.0.1"
        self.mongo_port = 27017
        self.mongo_db = ["admin", ]
        self.mongo_user = None
        self.mongo_password = None
        self.__conn = None
        self.__dbnames = None
        self.__metrics = {}

    def connect(self):
        """Connect to MongoDB"""
        if self.__conn is None:
            if self.mongo_user is None:
                try:
                    self.__conn = MongoClient('mongodb://%s:%s' %
                                              (self.mongo_host,
                                               self.mongo_port))
                except errors.PyMongoError as py_mongo_error:
                    print('Error in MongoDB connection: %s' %
                          str(py_mongo_error))
            else:
                try:
                    self.__conn = MongoClient('mongodb://%s:%s@%s:%s' %
                                              (self.mongo_user,
                                               self.mongo_password,
                                               self.mongo_host,
                                               self.mongo_port))
                except errors.PyMongoError as py_mongo_error:
                    print('Error in MongoDB connection: %s' %
                          str(py_mongo_error))

    def print_json(self):
        """print out all """
        metrics = self.__metrics
        print(json.dumps(metrics, indent=2, sort_keys=True))

    def get_discovery(self):
        """get discover"""
        if self.__conn is None:
            self.connect()
        db_handler = self.__conn[self.mongo_db[0]]
        try:
            rs_status = db_handler.command('replSetGetStatus')
        except errors.OperationFailure as exc:
            # if exc["codeName"] == 'NoReplicationEnabled':
            if exc.code == 76:
                return
        data = {"data": [{'{#APPLICATION}': 'rs'}]}
        print(json.dumps(data, indent=2, sort_keys=True))

    def get_rs_status(self):
        """get a list of DB names"""
        if self.__conn is None:
            self.connect()
        db_handler = self.__conn[self.mongo_db[0]]
        try:
            rs_status = db_handler.command('replSetGetStatus')
        except errors.OperationFailure as exc:
            # if exc["codeName"] == 'NoReplicationEnabled':
            if exc.code == 76:
                return
        # print(rs_status["members"])
        ok = rs_status["ok"]
        # 节点数
        if ok:
            # dict_metrics['value'] = len(rs_status["members"])
            self.__metrics['mongodb_rs_members'] = len(rs_status["members"])
        else:
            self.__metrics['mongodb_rs_members'] = 0
        # 延迟秒数
        primary_optime = None
        secondary_optime = None
        for member in rs_status["members"]:
            # print(member["_id"])
            if member["stateStr"] == "PRIMARY":
                primary_optime = member["optimeDate"]
            if member["stateStr"] == "SECONDARY":
                if secondary_optime is None:
                    secondary_optime = member["optimeDate"]
                else:
                    if time.mktime(member["optimeDate"].timetuple()) < time.mktime(secondary_optime.timetuple()):
                        # 保留一个最小的值
                        secondary_optime = member["optimeDate"]
        if primary_optime and secondary_optime:
            behind_seconds = time.mktime(primary_optime.timetuple()) - time.mktime(secondary_optime.timetuple())
        self.__metrics['mongodb_rs_behind_seconds'] = behind_seconds


    def close(self):
        """close connection to mongo"""
        if self.__conn is not None:
            self.__conn.close()

#######################################################################################################
### MAIN
if __name__ == '__main__':

    if ACTION == "discover":
        mongodb = MongoDB()
        mongodb.get_discovery()
        mongodb.close()
    elif ACTION == "rs":
        mongodb = MongoDB()
        mongodb.get_rs_status()
        mongodb.print_json()
        mongodb.close()
    else:
       usage()
